"""Interface specification for worker and worker manager.

This module defines the contract between the worker and worker manager components.
It specifies the workspace structure, expected file formats, and interaction patterns.
"""

from dataclasses import dataclass
import abc
from typing import Dict, List, Optional, Any, Union, Type
from tasks.model import TaskStatus, TaskSpec


@dataclass
class WorkspaceSpec:
    """Specification for the task workspace structure."""
    
    task_hash: str
    """Unique identifier for the task."""
    
    root_dir: str
    """Root directory where all task workspaces are stored."""
    
    @property
    def task_dir(self) -> str:
        """Full path to the task directory."""
        import os
        return os.path.join(self.root_dir, self.task_hash)
    
    @property
    def task_spec_path(self) -> str:
        """Path to the task specification file."""
        import os
        from tasks.worker.constants import FILE_TASK_SPEC
        return os.path.join(self.task_dir, FILE_TASK_SPEC)
    
    @property
    def repo_dir(self) -> str:
        """Path to the git repository directory."""
        import os
        from tasks.worker.constants import DIR_ROOT
        return os.path.join(self.task_dir, DIR_ROOT)
    
    @property
    def work_dir(self) -> str:
        """Path to the working directory."""
        import os
        from tasks.worker.constants import DIR_WORK
        return os.path.join(self.task_dir, DIR_WORK)
    
    @property
    def params_dir(self) -> str:
        """Path to the parameters directory."""
        import os
        from tasks.worker.constants import DIR_PARAMS
        return os.path.join(self.task_dir, DIR_PARAMS)
    
    @property
    def outputs_dir(self) -> str:
        """Path to the outputs directory."""
        import os
        from tasks.worker.constants import DIR_OUTPUTS
        return os.path.join(self.task_dir, DIR_OUTPUTS)
    
    @property
    def logs_dir(self) -> str:
        """Path to the logs directory."""
        import os
        from tasks.worker.constants import DIR_LOGS
        return os.path.join(self.task_dir, DIR_LOGS)
    
    @property
    def static_dir(self) -> str:
        """Path to the static cache directory."""
        import os
        from tasks.worker.constants import DIR_WORK, DIR_STATIC
        return os.path.join(self.work_dir, DIR_STATIC)
    
    def param_dir(self, param_name: str) -> str:
        """Path to a specific parameter directory."""
        import os
        from tasks.worker.constants import DIR_PARAMS
        return os.path.join(self.task_dir, DIR_PARAMS, param_name)
    
    def param_file(self, param_name: str) -> str:
        """Path to a specific file parameter."""
        import os
        from tasks.worker.constants import FILE_PARAM
        return os.path.join(self.param_dir(param_name), FILE_PARAM)


def create_workspace_structure(workspace: WorkspaceSpec) -> None:
    """Create the directory structure for a task workspace.
    
    Args:
        workspace: Workspace specification
    """
    import os
    from tasks.worker.constants import (
        DIR_ROOT, DIR_WORK, DIR_PARAMS, DIR_OUTPUTS, DIR_LOGS
    )
    
    os.makedirs(workspace.task_dir, exist_ok=True)
    os.makedirs(os.path.join(workspace.task_dir, DIR_ROOT), exist_ok=True)
    os.makedirs(os.path.join(workspace.task_dir, DIR_WORK), exist_ok=True)
    os.makedirs(os.path.join(workspace.task_dir, DIR_PARAMS), exist_ok=True)
    os.makedirs(os.path.join(workspace.task_dir, DIR_OUTPUTS), exist_ok=True)
    os.makedirs(os.path.join(workspace.task_dir, DIR_LOGS), exist_ok=True)


def key_tag_parse(key: str) -> tuple[str, list[str]]:
    """Parse a parameter key with optional tags.
    
    Args:
        key: Parameter key in format "name:tag1:tag2:..."
        
    Returns:
        Tuple of (parameter_name, [tags])
    """
    if ":" in key:
        fields = key.split(":")
        k = fields[0]
        tags = fields[1:]
    else:
        k = key
        tags = []
    return k, tags


def file_name_for_param(index: int) -> str:
    """Generate a filename for a parameter or output at the given index.
    
    Args:
        index: Parameter index
        
    Returns:
        Formatted filename string
    """
    return f"{index:08d}"


class WorkerInterface(abc.ABC):
    """Interface specification for worker behavior.
    
    Workers are responsible for executing tasks in a prepared workspace.
    They are instantiated with a workspace and provide methods to start,
    monitor, and get the status of a task.
    """
    
    @abc.abstractmethod
    def __init__(self, workspace: WorkspaceSpec):
        """Initialize a worker for a specific workspace.
        
        Args:
            workspace: The workspace specification for the task
        """
        pass
    
    @abc.abstractmethod
    async def start(self) -> None:
        """Start the task execution.
        
        Returns:
            None
        
        Raises:
            Exception: If the task fails to start
        """
        pass
    
    @abc.abstractmethod
    async def get_status(self) -> TaskStatus:
        """Get the current status of the task.
        
        Returns:
            TaskStatus: The current status of the task
        """
        pass
    
    @abc.abstractmethod
    async def wait(self) -> int:
        """Wait for the task to complete and return the exit code.
        
        Returns:
            int: Exit code (0 for success, non-zero for failure)
        """
        pass
    
    @abc.abstractmethod
    async def cancel(self) -> None:
        """Cancel the running task.
        
        Returns:
            None
        """
        pass
    
    @property
    @abc.abstractmethod
    def returncode(self) -> Optional[int]:
        """Get the return code of the task if it has completed.
        
        Returns:
            Optional[int]: Return code if task has completed, None otherwise
        """
        pass


class ManagerInterface(abc.ABC):
    """Interface specification for worker manager behavior."""
    
    @abc.abstractmethod
    def prepare_workspace(self, task_spec: TaskSpec, workspace: WorkspaceSpec) -> None:
        """Prepare workspace for a task.
        
        Args:
            task_spec: Task specification
            workspace: Workspace specification
        """
        pass
    
    @abc.abstractmethod
    async def run_task(self, task: TaskSpec, worker_class: Type[WorkerInterface]) -> None:
        """Run a task using the specified worker class.
        
        Args:
            task: Task specification
            worker_class: Worker class to use for executing the task
        """
        pass 